import socket
import cv2 as cv
import numpy as np
import os
import shutil
import easygui as a

redlower = np.array([0, 43, 46])
redupper = np.array([10, 255, 255])
bluelower = np.array([100, 43, 46])
blueupper = np.array([124, 255, 255])
greenlower = np.array([37, 43, 46])
greenupper = np.array([77, 255, 255])


def ReceivePicture():
    address = ('127.0.0.1', 3000)
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.bind(address)
    s.listen(1)

    def recvall(sock, count):
        buf = b''
        while count:
            newbuf = sock.recv(count)
            if not newbuf: return None
            buf += newbuf
            count -= len(newbuf)
        return buf
    conn, addr = s.accept()
    print('connect from:' + str(addr))
    i = 1
    while 1:
        length = recvall(conn, 16)  # 获得图片文件的长度,16代表获取长度
        stringData = recvall(conn, int(length))  # 根据获得的文件长度，获取图片文件
        data = np.frombuffer(stringData, np.uint8)  # 将获取到的字符流数据转换成1维数组
        decimg = cv.imdecode(data, cv.IMREAD_COLOR)  # 将数组解码成图像
        cv.imwrite(r"./pics/pic" + str(i) + ".jpg", decimg)
        i += 1
        if (i == 250):
            s.close()
            break;


def showpro():
    choice = a.ccbox(msg='选择需要显示红色(1)还是绿色(2)。', title='红绿色盲辅助脚本', choices=('1', '2'))
    if (choice == 1):
        for i in range(1, 250):
            img = cv.imread(r"./pics/pic" + str(i) + ".jpg", 1)
            hsv = cv.cvtColor(img, cv.COLOR_BGR2HSV)
            maskG = cv.inRange(hsv, lowerb=greenlower, upperb=greenupper)
            maskR = cv.inRange(hsv, lowerb=redlower, upperb=redupper)
            maskRG = cv.bitwise_or(maskG, maskR)
            picpro = cv.bitwise_or(img, img, mask=maskR)
            out = cv.GaussianBlur(picpro, (3, 3), 1.3)  # σ较大，则生成的模板的各个系数相差就不是很大，比较类似均值模板，对图像的平滑效果比较明显。
            cv.imshow('videopro', out)
            cv.imshow('videorecv', img)
            c = cv.waitKey(35)
            if c == 27:
                break;
    else:
        for i in range(1, 250):
            img = cv.imread(r"./pics/pic" + str(i) + ".jpg", 1)
            hsv = cv.cvtColor(img, cv.COLOR_BGR2HSV)
            maskG = cv.inRange(hsv, lowerb=greenlower, upperb=greenupper)
            maskR = cv.inRange(hsv, lowerb=redlower, upperb=redupper)
            maskRG = cv.bitwise_or(maskG, maskR)
            picpro = cv.bitwise_or(img, img, mask=maskG)
            out = cv.GaussianBlur(picpro, (3, 3), 1.3)  # σ较大，则生成的模板的各个系数相差就不是很大，比较类似均值模板，对图像的平滑效果比较明显。
            cv.imshow('videopro', out)
            cv.imshow('videorecv', img)
            c = cv.waitKey(35)
            if c == 27:
                break;


if __name__ == '__main__':
    ReceivePicture()
    showpro()
    shutil.rmtree(r"./picturedata")
    os.mkdir(r"./picturedata")